package com.spark.sql

import org.apache.spark.sql.expressions.UserDefinedFunction
import org.apache.spark.sql.{DataFrame, SparkSession}

object Demo13UDF {
  def main(args: Array[String]): Unit = {
    val spark: SparkSession = SparkSession
      .builder()
      .master("local")
      .appName("Demo13UDF")
      .config("spark.sql.shuffle.partitions", 1)
      .getOrCreate()

    spark.sparkContext.setLogLevel("error")

    import org.apache.spark.sql.functions._
    import spark.implicits._

    /**
     * 1、DSL自定义函数
     * 在自定义函数中可以编写scala代码
     */
    val str_len: UserDefinedFunction = udf((str: String) => {
      str.length
    })

    val studentDF: DataFrame = spark
      .read
      .format("json")
      .load("data/students.json")

    //使用自定义函数
//    studentDF
//      .select(str_len($"clazz"))
//      .show()

    /**
     * 2、SQL 自定义函数
     *
     */
    //在sql中使用自定义函数需要注册
    spark.udf.register("str_len", str_len)
    studentDF.createOrReplaceTempView("student")

    spark.sql(
      """
        |select *, str_len(clazz) from student
        |""".stripMargin)
      .show()

  }

}
